"""
AI服务模块 - 负责与DeepSeek API的交互
"""

import json
from openai import OpenAI
from tools import available_tools, tool_functions
import config

class AIService:
    def __init__(self, api_key, base_url):
        self.api_key = api_key
        self.base_url = base_url
        self.client = OpenAI(api_key=api_key, base_url=base_url)
    
    def chat(self, user_message, custom_tools=None):
        """
        处理用户消息，调用AI API并处理响应
        
        Args:
            user_message: 用户消息文本
            custom_tools: 可选，自定义工具集列表
        """
        # 确定使用的工具集
        tools_to_use = custom_tools if custom_tools else available_tools
        
        # 调用API
        response = self.client.chat.completions.create(
            model=config.MODEL_NAME,
            messages=[
                {"role": "system", "content": config.SYSTEM_PROMPT},
                {"role": "user", "content": user_message},
            ],
            tools=tools_to_use,
            stream=False
        )
        
        # 获取AI响应
        ai_message = response.choices[0].message
        
        # 检查是否有工具调用
        if hasattr(ai_message, 'tool_calls') and ai_message.tool_calls:
            return self._handle_tool_calls(user_message, ai_message, tools_to_use)
        
        # 如果没有工具调用，直接返回响应
        return {"response": ai_message.content, "has_tool_calls": False}
    
    def _handle_tool_calls(self, user_message, ai_message, tools_to_use):
        """
        处理工具调用
        """
        tool_responses = []
        
        # 处理每个工具调用
        for tool_call in ai_message.tool_calls:
            function_name = tool_call.function.name
            function_args = json.loads(tool_call.function.arguments)
            
            # 执行工具函数
            if function_name in tool_functions:
                function_result = tool_functions[function_name](**function_args)
                tool_responses.append({
                    "tool_call_id": tool_call.id,
                    "function_name": function_name,
                    "result": function_result
                })
        
        # 获取工具调用结果后的后续响应
        messages = [
            {"role": "system", "content": config.SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
            ai_message.model_dump(),
        ]
        
        # 添加工具响应
        for tool_response in tool_responses:
            messages.append({
                "role": "tool",
                "tool_call_id": tool_response["tool_call_id"],
                "content": json.dumps(tool_response["result"], ensure_ascii=False)
            })
        
        # 再次调用API获取最终响应
        final_response = self.client.chat.completions.create(
            model=config.MODEL_NAME,
            messages=messages,
            stream=False
        )
        
        return {
            "response": final_response.choices[0].message.content,
            "has_tool_calls": True,
            "tool_calls": [{
                "function_name": tc["function_name"],
                "result": tc["result"]
            } for tc in tool_responses]
        } 